import scala.collection.mutable.ArrayBuffer
import scala.util.Random
import java.awt.image.{BufferedImage}
import com.sksamuel.scrimage.ImmutableImage
import com.sksamuel.scrimage.color.RGBColor
import com.sksamuel.scrimage.nio.PngWriter

def ray_color(r: Ray, world: Hittable, depth: Int): Vec3 =
  if depth <=0 then
    return Vec3(0, 0, 0)

  world.hit(r, 0.001, Double.MaxValue) match
    case None => 
      val unit_direction = r.direction.unit_vector
      val t = 0.5 * (unit_direction.y + 1)
      Vec3(1, 1, 1) * (1 - t) + Vec3(0.5, 0.7, 1.0) * t
    case Some(rec) =>
      rec.material.scatter(r, rec) match
        case Some((scattered, attenuation)) =>
          ray_color(scattered, world, depth - 1) * attenuation
        case None => Vec3(0, 0, 0)


@main def main(): Unit =
  val t0 = System.currentTimeMillis
  val aspect_ratio = 16.0 / 9.0
  val image_width = 400
  val image_height = (image_width / aspect_ratio).toInt
  val samples_per_pixel = 100
  val max_depth = 50

  val rng = Random()

  val world = HittableList()
  world.add(Sphere(
    Vec3(0, -1000, 0),
    1000.0,
    Lambertian(Vec3(0.5, 0.5, 0.5))
  ))

  for
    a <- -11 until 11
    b <- -11 until 11
  do
    val choose_mat = rng.nextDouble
    val center = Vec3(
      0.9 * rng.nextDouble + a,
      0.2,
      0.9 * rng.nextDouble + b
    )
    if (center - Vec3(4, 0.2, 0)).length > 0.9 then
      if choose_mat < 0.8 then
        world.add(Sphere(
          center, 0.2,
          Lambertian(Vec3(
            rng.nextDouble * rng.nextDouble,
            rng.nextDouble * rng.nextDouble,
            rng.nextDouble * rng.nextDouble,
          ))
        ))
      else if choose_mat < 0.95 then
        world.add(Sphere(
          center, 0.2,
          Metal(Vec3(
            0.5 * (1.0 + rng.nextDouble),
            0.5 * (1.0 + rng.nextDouble),
            0.5 * (1.0 + rng.nextDouble),
          ), 0.5 * rng.nextDouble)
        ))
      else
        world.add(Sphere(
          center, 0.2,
          Dielectric(1.5)
        ))
  world.add(Sphere(
    Vec3(0, 1, 0), 1.0,
    Dielectric(1.5)
  ))
  world.add(Sphere(
    Vec3(-4, 1, 0), 1.0,
    Lambertian(Vec3(0.4, 0.2, 0.1))
  ))
  world.add(Sphere(
    Vec3(4, 1, 0),
    1.0,
    Metal(Vec3(0.7, 0.6, 0.5), 0)
  ))

  val lookfrom = Vec3(13, 2, 3)
  val lookat = Vec3(0, 0, 0)
  val vup = Vec3(0, 1, 0)
  val dist_to_focus = 10
  val aperture = 0.1
  val cam = Camera(lookfrom, lookat, vup, 20.0, aspect_ratio, aperture, dist_to_focus)
  var image = ImmutableImage.create(image_width, image_height, BufferedImage.TYPE_INT_RGB)

  image = image.map{pixel => 
    val i = pixel.x
    val j = image_height - pixel.y - 1

    // if i == 0 then
    //   println(s"Scanrows remaining: $j")

    var pixel_color = Vec3(0, 0, 0)

    for _ <- 0 until samples_per_pixel do
      val u = (1.0 * i + rng.nextDouble) / (image_width -1)
      val v = (1.0 * j + rng.nextDouble) / (image_height - 1)
      val r = cam.get_ray(u, v)
      val color = ray_color(r, world, max_depth)
      pixel_color += color

    pixel_color.to_color(samples_per_pixel)
  }

  image.output(PngWriter(), "./out/a.png")

  println(s"finish in ${(System.currentTimeMillis - t0) / 1000.0} s")

